/*
 * main.cpp - raytracing program main executable
 * Wingfai Chan, Daniel Kotowski
 * CIS 4450
 *
 */

#ifdef _WIN32
#  define WINDOWS_LEAN_AND_MEAN
#  define NOMINMAX
#  include <windows.h>
#pragma warning(disable:4996)
#endif


//OpenGL
#include <GL/glew.h>
#include <GL/glut.h>

//CUDA
#include <cuda_runtime_api.h>
#include <cutil.h>
#include <cutil_inline.h>
#include <cutil_gl_inline.h>
#include <cutil_gl_error.h>
#include <cuda_gl_interop.h>
#include <vector_types.h>
#include <vector_functions.h>
#include <cutil_math.h>

#include <AntTweakBar.h> 
#include <stdio.h>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <math.h>

using namespace std;

//Structures
struct TriangleFace 
{
	int v[3]; // vertex indices
	int n[3]; // vertex normals
};

struct TriangleMesh
{
	vector<float3> verts;
	vector<float3> normals;
	vector<TriangleFace> faces;
	//float3 bounding_box[2];
	float4 colour;
};

//Variables

int horizontal_res = 800;
int vertical_res = 600;

//camera variables
float3 camera_pos;
float3 a, b, c;
float camera_rot = 1;
float camera_dist = 25;
float camera_height = 25;
float time_delta;
bool animate = true;

float3 light_pos = make_float3(0.0, 50.0, 0.0);

GLuint pixelbuffer;
GLuint display_texture;
vector<float4> triangles;
vector<float4> bounding_spheres;

//Function prototypes
void createPBO(GLuint *, int, int);
void createTexture(GLuint *, unsigned int, unsigned int);
void deletePBO(GLuint *);
void deleteTexture(GLuint *);

void display();

void drawSceneTexture();
void renderScene();

void initCUDA();
void initGUI();
void initOpenGL(int, char **);

void updateCamera();
void loadObj(const string, TriangleMesh &, float3, float4, float &);
void addObjToScene(const string, float3, float4);

extern "C" void launchRayTrace(unsigned int*, int, int, int, float3, float3, float3, float3, float3, float3);
extern "C" void bindTriangleTexture(float *, size_t);
extern "C" void bindBoundingTexture(float *, size_t);
extern "C" float dist(float3, float3);

/*
 * This function performs some opengl initialization calls
 *
 */
void initOpenGL(int *argc, char **argv)
{
	glutInit(argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB);
    glutInitWindowSize(horizontal_res, vertical_res); //Set the window size
	
	//we don't need depth testing because we are drawing the scene manually
	glDisable(GL_DEPTH_TEST);
	glClear(GL_COLOR_BUFFER_BIT);

	glViewport(0, 0, horizontal_res, vertical_res);

	glutCreateWindow("CIS4300 Raytracing with CUDA");

	//check that support for the extensions are available
	glewInit();
	if (!glewIsSupported("GL_VERSION_2_0 " "GL_ARB_pixel_buffer_object " "GL_EXT_framebuffer_object ")) 
	{
		printf("ERROR: Support for necessary OpenGL extensions missing.\n");
		exit(0);
	}
}

/*
 * This function sets up the GUI menu.
 * Uses AntTweakBar http://www.antisphere.com/Wiki/tools:anttweakbar
 *
 */
void initGUI()
{
	TwBar *bar;

	if (!TwInit(TW_OPENGL, NULL))
	{
		printf("GUI failed to initialize.\n");
		exit(0);
	}

	TwWindowSize(horizontal_res, vertical_res);

	bar = TwNewBar("Controls");

	TwAddVarRW(bar, "Light position - x", TW_TYPE_FLOAT, &light_pos.x, 
		" min= -100.0 max= 100.0 step=0.1 group='Light Position'");
	TwAddVarRW(bar, "Light position - y", TW_TYPE_FLOAT, &light_pos.y, 
		" min= -100.0 max= 100.0 step=0.1 group='Light Position'");
	TwAddVarRW(bar, "Light position - z", TW_TYPE_FLOAT, &light_pos.z, 
		" min= -100.0 max= 100.0 step=0.1 group='Light Position'");

	TwAddVarRW(bar, "Camera rotation", TW_TYPE_FLOAT, &camera_rot, 
		" min= -5 max= 5 step=0.1 group='Camera'");
	TwAddVarRW(bar, "Camera height", TW_TYPE_FLOAT, &camera_height, 
		" min= -100 max= 100 step=0.1 group='Camera'");
	TwAddVarRW(bar, "Camera distance", TW_TYPE_FLOAT, &camera_dist, 
		" min= 0 max= 100 step=0.1 group='Camera'");

}


/*
 * This function binds a openGL buffer to the CUDA device so that we can pass data from
 * the device to openGL; it also loads the obj data and binds it to the CUDA device
 */
void initCUDA()
{
	//create the open gl textures and buffers 
	cudaGLSetGLDevice(cutGetMaxGflopsDeviceId());
	createPBO(&pixelbuffer, horizontal_res, vertical_res);
	createTexture(&display_texture, horizontal_res, vertical_res);

	
	//load the objs
	addObjToScene("obj/8cubes1plane_Scene.obj", make_float3(0, 0, 0), make_float4(100, 50, 255, 0.5));
	//addObjToScene("obj/droideka.obj", make_float3(0, 0, 0), make_float4(128, 128, 128, 0.99));
	//addObjToScene("obj/enzo_studio.obj", make_float3(0, 0, 0), make_float4(128, 128, 128, 0));
	//addObjToScene("obj/monkey.obj", make_float3(0, 0, 0), make_float4(128, 128, 128, 1));
	addObjToScene("obj/sphere.obj", make_float3(0, 8, 0), make_float4(128, 128, 128, 1));
	//addObjToScene("obj/sphere.obj", make_float3(40, 0, 0), make_float4(255, 125, 0, 0.99));
	//addObjToScene("obj/box.obj", make_float3(0, -100, 0), make_float4(128, 128, 128, 0));
	
	printf("Total number of triangles: %d\n", triangles.size()/7);
	float *device_triangles;
	size_t size = triangles.size() * sizeof(float4);

	//malloc and copy the triangles to device memory
	cutilSafeCall(cudaMalloc((void**)&device_triangles, size));
	cutilSafeCall(cudaMemcpy(device_triangles, &triangles[0], size, cudaMemcpyHostToDevice));

	//bind the triangles to a texture on the device
	bindTriangleTexture(device_triangles, size);


	//allocate and bind bounding spheres
	float *device_bounding;
	size = bounding_spheres.size() * sizeof(float4);

	//malloc and copy the bounding spheres to device memory
	cutilSafeCall(cudaMalloc((void**)&device_bounding, size));
	cutilSafeCall(cudaMemcpy(device_bounding, &bounding_spheres[0], size, cudaMemcpyHostToDevice));

	//bind the bounding volumes
	bindBoundingTexture(device_bounding, size);
	printf("Total number of objs: %d\n", bounding_spheres.size()/2);
}

/*
 * Start the kernel to raytrace the scene and download the data into the buffer so it may be displayed.
 *
 */
void renderScene()
{
	unsigned int *out_data;
	cutilSafeCall(cudaGLMapBufferObject((void**)&out_data, pixelbuffer));

	launchRayTrace(out_data, horizontal_res, vertical_res, bounding_spheres.size(), camera_pos, a, b, c, light_pos, make_float3(255,255,255));
	cudaThreadSynchronize();

	cutilSafeCall(cudaGLUnmapBufferObject(pixelbuffer));

	// download texture from destination PBO
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, pixelbuffer);
	glBindTexture(GL_TEXTURE_2D, display_texture);
	glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, horizontal_res, vertical_res, GL_BGRA, GL_UNSIGNED_BYTE, NULL);
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, 0);

	CUT_CHECK_ERROR_GL();
}

void updateCamera()
{
	camera_pos = make_float3(cos(camera_rot)*camera_dist, camera_height, -sin(camera_rot)*camera_dist);
	float3 camera_dir = normalize(-camera_pos);
	
	float3 cam_up  = make_float3(0,1,0);
	float3 cam_right = cross(camera_dir,cam_up);
	cam_right = normalize(cam_right);

	cam_up = -cross(camera_dir,cam_right);
	cam_up = normalize(cam_up);
	
	float FOV = 60.0f;
	float theta = (FOV*3.1415*0.5) / 180.0f;
	float half_width = tanf(theta);
	float aspect = (float)horizontal_res / (float)vertical_res;

	float u0 = -half_width * aspect;
	float v0 = -half_width;
	float u1 =  half_width * aspect;
	float v1 =  half_width;
	float dist_to_image = 1;

	a = (u1-u0)*cam_right;
	b = (v1-v0)*cam_up;
	c = camera_pos + u0*cam_right + v0*cam_up + dist_to_image*camera_dir;
		
	if (animate)
		camera_rot += 0.25 * time_delta;
}

void addObjToScene(const string filename, float3 offset, float4 colour)
{
	TriangleMesh mesh;
	unsigned int i, start_index, end_index;
	float3 vertex0, vertex1, vertex2;
	float3 normal0, normal1, normal2;
	float radius;
	
	//load the objs
	loadObj(filename, mesh, offset, colour, radius);
	
	start_index = triangles.size();
	//add to the master list of triangles
	for (i=0; i<mesh.faces.size(); i++)
	{
		//retrieve the vertex coordinates of the face
		vertex0 = mesh.verts[mesh.faces[i].v[0]-1];
		vertex1 = mesh.verts[mesh.faces[i].v[1]-1];
		vertex2 = mesh.verts[mesh.faces[i].v[2]-1];
		
		//retrieve the vertex normals of the face
		normal0 = mesh.normals[mesh.faces[i].n[0]-1];
		normal1 = mesh.normals[mesh.faces[i].n[1]-1];
		normal2 = mesh.normals[mesh.faces[i].n[2]-1];

		//store one vertex and two edges into the list of triangles
		//the reasoning for this is for the triangle/ray intersection algorithm rayTriangleIntersect() in which
		//the edge conversions can be done as a pre process
		triangles.push_back(make_float4(vertex0));
		triangles.push_back(make_float4(vertex1-vertex0)); //edge from vertex 1 to vertex 0
		triangles.push_back(make_float4(vertex2-vertex0)); //edge from vertex 2 to vertex 0
		triangles.push_back(mesh.colour);

		triangles.push_back(make_float4(normal0));
		triangles.push_back(make_float4(normal1));
		triangles.push_back(make_float4(normal2));
	}
	end_index = triangles.size()-1;

	//add the bounding sphere
	bounding_spheres.push_back(make_float4(offset, radius));
	bounding_spheres.push_back(make_float4((float)start_index, (float)end_index, 0.0, 0.0));
}

/*
 * load a simple obj file without normals or tex-coords
 * Note: this function was taken from the tutorial found here
 * http://cg.alexandra.dk/2009/08/10/triers-cuda-ray-tracing-tutorial/
 * and modified
 */
void loadObj(const string filename, TriangleMesh &mesh, float3 offset, float4 colour, float &radius)
{

	ifstream in(filename.c_str());

	if(!in.good())
	{
		cout  << "ERROR: loading obj:(" << filename << ") file is not good" << "\n";
		exit(0);
	}

	char buffer[256], str[255];
	float f1,f2,f3;
	radius = 0.0;

	while(!in.getline(buffer,255).eof())
	{
		buffer[255]='\0';

		sscanf_s(buffer,"%s",str,255);

		// reading a vertex
		if (buffer[0]=='v' && (buffer[1]==' '  || buffer[1]==32) )
		{
			if ( sscanf(buffer,"v %f %f %f",&f1,&f2,&f3)==3)
			{
				mesh.verts.push_back(make_float3(f1+offset.x,f2+offset.y,f3+offset.z));

				//find the radius of a sphere that will enclose the obj
				if (dist(mesh.verts.back(), offset) > radius)
					radius = dist(mesh.verts.back(), offset);
			}
			else
			{
				cout << "ERROR: vertex not in wanted format in OBJLoader" << "\n";
				exit(-1);
			}
		}
		
		//read in the vertex normals
		if (buffer[0]=='v' && buffer[1] == 'n' && (buffer[2] == ' ' || buffer[2] ==32))
		{
			if (sscanf(buffer, "vn %f %f %f", &f1, &f2, &f3) == 3)
				mesh.normals.push_back(make_float3(f1, f2, f3));
			else
			{
				printf("ERROR: vertex normals not in wanted format\n");
				exit(-1);
			}
		}

		// reading FaceMtls 
		else if (buffer[0]=='f' && (buffer[1]==' ' || buffer[1]==32) )
		{
			TriangleFace f;
			int nt = sscanf(buffer,"f %d//%d %d//%d %d//%d",&f.v[0], &f.n[0], &f.v[1], &f.n[1], &f.v[2],&f.n[2]);
			if( nt!=6 )
			{
				cout << "ERROR: I don't know the format of that FaceMtl" << "\n";
				exit(-1);
			}
			
			mesh.faces.push_back(f);
		}
	}
	mesh.colour = colour;
}



//=======================================================================================
//openGL buffer and texture functions from post processing example from the CUDA SDK

/*
 * Create a pixel buffer object.
 */
void createPBO(GLuint* pbo, int width, int height)
{
    // set up vertex data parameter
    unsigned int num_texels = width * height;
    unsigned int num_values = num_texels * 4;
    unsigned int size_tex_data = sizeof(GLubyte) * num_values;

    void *data = malloc(size_tex_data);

    // create buffer object
    glGenBuffers(1, pbo);
    glBindBuffer(GL_ARRAY_BUFFER, *pbo);
    glBufferData(GL_ARRAY_BUFFER, size_tex_data, data, GL_DYNAMIC_DRAW);
    free(data);

    glBindBuffer(GL_ARRAY_BUFFER, 0);

    // register this buffer object with CUDA
    cutilSafeCall(cudaGLRegisterBufferObject(*pbo));

    CUT_CHECK_ERROR_GL();
}

/*
 * Delete a pixel buffer object.
 */
void deletePBO(GLuint* pbo)
{
    // unregister this buffer object with CUDA
    cutilSafeCall(cudaGLUnregisterBufferObject(*pbo));

    glBindBuffer(GL_ARRAY_BUFFER, *pbo);
    glDeleteBuffers(1, pbo);
    CUT_CHECK_ERROR_GL();

    *pbo = 0;
}

/*
 * Create and bind a texture.  
 */
void createTexture(GLuint* tex_name, unsigned int size_x, unsigned int size_y)
{
    // create a texture
    glGenTextures(1, tex_name);
    glBindTexture(GL_TEXTURE_2D, *tex_name);

    // set basic parameters
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);

    // buffer data
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, size_x, size_y, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);

    CUT_CHECK_ERROR_GL();
}

/*
 * Unbind a texture.
 */
void deleteTexture(GLuint* tex)
{
    glDeleteTextures(1, tex);
    CUT_CHECK_ERROR_GL();

    *tex = 0;
}

/*
 * Takes the data from the pbo and creates a screen sized 2D quad with the texture.
 *
 */
void drawSceneTexture()
{
	glDisable(GL_DEPTH_TEST);
	glDisable(GL_LIGHTING);
	glEnable(GL_TEXTURE_2D);
	glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE);

	glMatrixMode(GL_PROJECTION);
	glPushMatrix();
	glLoadIdentity();
	glOrtho(-1.0, 1.0, -1.0, 1.0, -1.0, 1.0);

	glMatrixMode( GL_MODELVIEW);
	glLoadIdentity();

	glViewport(0, 0, horizontal_res, vertical_res);

	//draw the quad
	glBegin(GL_QUADS);
	glTexCoord2f(0.0, 0.0); glVertex3f(-1.0, -1.0, 0.5);
	glTexCoord2f(1.0, 0.0); glVertex3f(1.0, -1.0, 0.5);
	glTexCoord2f(1.0, 1.0); glVertex3f(1.0, 1.0, 0.5);
	glTexCoord2f(0.0, 1.0); glVertex3f(-1.0, 1.0, 0.5);
	glEnd();

	glMatrixMode(GL_PROJECTION);
	glPopMatrix();

	glDisable(GL_TEXTURE_2D);
	CUT_CHECK_ERROR_GL();
}
//=======================================================================================








//=======================================================================================
//OPENGL CALLBACKS

/*
 * Display callback
 */
void display()
{
	//update the delta time for animation
	static int lastFrameTime = 0;
	if (lastFrameTime == 0)
	{
		lastFrameTime = glutGet(GLUT_ELAPSED_TIME);
	}

	int now = glutGet(GLUT_ELAPSED_TIME);
	int elapsedMilliseconds = now - lastFrameTime;
	//printf("%d\n", elapsedMilliseconds);
	time_delta = elapsedMilliseconds / 1000.0f;
	lastFrameTime = now;

	updateCamera();

	renderScene();
	drawSceneTexture();

	//draw the menu bar
	TwDraw();

	glutSwapBuffers();
	glutPostRedisplay();
}

/*
 * Keyboard callback
 */
void keyboard(unsigned char key, int x, int y)
{
	switch (key)
	{
		case ' ':
			animate = !animate;
			break;
	}
}
//=======================================================================================


int main(int argc, char **argv)
{
	initOpenGL(&argc, argv);
	initGUI();
	initCUDA();

	glutDisplayFunc(display);
	glutKeyboardFunc(keyboard);

	// - Directly redirect GLUT mouse button events to AntTweakBar
	glutMouseFunc((GLUTmousebuttonfun)TwEventMouseButtonGLUT);
	// - Directly redirect GLUT mouse motion events to AntTweakBar
	glutMotionFunc((GLUTmousemotionfun)TwEventMouseMotionGLUT);
	// - Directly redirect GLUT mouse "passive" motion events to AntTweakBar (same as MouseMotion)
	glutPassiveMotionFunc((GLUTmousemotionfun)TwEventMouseMotionGLUT);


	//renderScene();
	glutMainLoop();
	cudaThreadExit();
	
	cutilExit(argc, argv);

	return (0);
}